module Trie (Trie(..),
             empty,
             isEmpty,
             insert,
             lkup,
             remove,
             mapT,
             mapTK,
             filterT,
             filterTK,
             foldT,
             foldTK,
             insertTrie,
             fromList,
             asList,
             singleton,
             update,
             restrict,
             removeSet,
             pointwise,
             pointwiseRan,
             pointwiseDom,
             mutually,
             mutuallyRan,
             mutuallyDom) where

import AList
import qualified Data.Set as S


-- Basic Trie functionality (problem set 1).

data Trie a = Node (Maybe a) (AList (Trie a)) deriving (Eq, Ord)


isEmpty :: Trie a -> Bool
isEmpty (Node Nothing ts) = isNil ts
isEmpty _                 = False


empty :: Trie a
empty = Node Nothing nil


insert :: String -> a -> Trie a -> Trie a
insert "" v (Node x ts)     = Node (Just v) ts
insert (c:cs) v (Node x ts) = case get c ts of
                                Nothing -> let t' = insert cs v empty
                                           in Node x (put c t' ts)
                                Just t  -> let t' = insert cs v t
                                           in Node x (put c t' ts)


lkup :: String -> Trie a -> Maybe a
lkup "" (Node x ts)     = x
lkup (c:cs) (Node x ts) = case get c ts of
                            Nothing -> Nothing
                            Just t  -> lkup cs t


remove :: String -> Trie a -> Trie a
remove "" (Node x ts)     = Node Nothing ts
remove (c:cs) (Node x ts) = case get c ts of
                              Nothing -> Node x ts
                              Just t  -> let t' = remove cs t
                                         in if isEmpty t'
                                            then Node x (del c ts)
                                            else Node x (put c t' ts)


instance Show a => Show (Trie a) where
    show t | isEmpty t = "{}"
    show t | otherwise = "{" ++ (tail $ foldTK (\ k v s -> s ++ "," ++ k ++ "->" ++ show v) "" t) ++ "}"


-- Higher order functions for Tries (problem set 2)

mapT :: (a -> b) -> Trie a -> Trie b
mapT f = mapTK (\ _ v -> f v)


mapTK :: (String -> a -> b) -> Trie a -> Trie b
mapTK f (Node v ts) = mapKV' "" f (Node v ts)
    where
      mapKV' pre f (Node v ts) = let v'  = case v of
                                             Nothing -> Nothing
                                             Just w  -> Just (f pre w)
                                     ts' = map (\ (k, t) -> (k, mapKV' (pre ++ [k]) f t)) ts
                                 in Node v' ts'


filterT :: (a -> Bool) -> Trie a -> Trie a
filterT p = filterTK (\ _ v -> p v)


filterTK :: (String -> a -> Bool) -> Trie a -> Trie a
filterTK p t = filterTK' "" p t
    where
      filterTK' pre p (Node v ts) = let v'  = case v of
                                                 Nothing -> Nothing
                                                 Just w  -> if p pre w then v else Nothing
                                        ts'  = map (\ (k, t) -> (k, filterTK' (pre ++ [k]) p t)) ts
                                        ts'' = filter (not . isEmpty . snd) ts'
                                    in Node v' ts''


foldT :: (a -> b -> b) -> b -> Trie a -> b
foldT f = foldTK (\ _ v a -> f v a)


foldTK :: (String -> a -> b -> b) -> b -> Trie a -> b
foldTK f a = foldTK' "" f a
    where
      foldTK' pre f a (Node Nothing ts)  = foldr (\ (k, t) r -> foldTK' (pre ++ [k]) f r t) a ts
      foldTK' pre f a (Node (Just v) ts) = let tfold = foldr (\ (k, t) r -> foldTK' (pre ++ [k]) f r t) a ts
                                           in f pre v tfold


insertTrie :: Trie a -> Trie a -> Trie a
insertTrie t1 t2 = foldTK insert t1 t2


asList :: Trie a -> [(String, a)]
asList t = foldTK (\ k v l -> (k, v):l) [] t


fromList :: [(String, a)] -> Trie a
fromList = foldr (\ (k, v) t -> insert k v t) empty


-- Map functions for tries (problem set 6)
-- % begin_map_funs % --
member :: String -> Trie a  -> Bool
member k t = case lkup k t of { Nothing -> False; Just _ -> True }

singleton :: String -> a -> Trie a
singleton k a = insert k a empty

dom :: Trie a -> S.Set String
dom = foldTK (\ k _ s -> S.insert k s) S.empty

ran :: Ord a => Trie a -> S.Set a
ran = foldT S.insert S.empty

update :: Trie a -> Trie a -> Trie a
update = insertTrie

removeSet :: S.Set String -> Trie a -> Trie a
removeSet s t = S.fold remove t s

restrict  :: Trie a -> S.Set String -> Trie a
restrict t s = filterTK (\ k _ -> S.member k s) t
-- % end_map_funs % --


-- Predicates on Tries (problem set 6).
-- % begin_predicates % --
pointwise :: ((String, a) -> Bool) -> Trie a -> Bool
pointwise p = foldT (&&) True . mapTK (curry p)


pointwiseRan :: (a -> Bool) -> Trie a -> Bool
pointwiseRan p = pointwise (p . snd)


pointwiseDom :: (String -> Bool) -> Trie a -> Bool
pointwiseDom p = pointwise (p . fst)


mutually :: ((String, a) -> (String, a) -> Bool) -> Trie a -> Bool
mutually p = mutually' . asList
    where
      mutually' []        = True
      mutually' (kv1:kvs) = foldr (\ kv2 r -> r && checkP kv1 kv2) True kvs
                            && mutually' kvs
      checkP kv1 kv2      = p kv1 kv2 && p kv2 kv1


mutuallyRan :: (a -> a -> Bool) -> Trie a -> Bool
mutuallyRan p = mutually (\ (_, v1) (_, v2) -> p v1 v2)


mutuallyDom :: (String -> String -> Bool) -> Trie a -> Bool
mutuallyDom p = mutually (\ (k1, _) (k2, _) -> p k1 k2)
-- % end_predicates % --
